"""
存储模块 - 负责将爬取的数据保存为结构化文件
"""

import os
import json
import time
import logging
from datetime import datetime
from urllib.parse import urlparse

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('logs/crawler.log'),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger('storage')

class Storage:
    """
    数据存储类
    
    负责:
    1. 将爬取的数据保存为JSON文件
    2. 按域名和日期组织文件结构
    3. 提供数据读取功能
    """
    
    def __init__(self, config):
        """
        初始化存储模块
        
        参数:
            config: 配置对象
        """
        self.base_path = config.get('DEFAULT', 'save_path', fallback='./data')
        
        # 确保存储目录存在
        if not os.path.exists(self.base_path):
            os.makedirs(self.base_path)
            logger.info(f"创建存储目录: {self.base_path}")
        
        logger.info(f"存储模块初始化完成: 数据将保存到 {self.base_path}")
    
    def get_save_path(self, url):
        """
        根据URL生成保存路径
        
        参数:
            url: 源URL
        
        返回:
            元组: (目录路径, 文件名)
        """
        # 解析域名
        domain = urlparse(url).netloc
        
        # 获取当前日期
        date_str = datetime.now().strftime('%Y%m%d')
        
        # 构建目录路径
        dir_path = os.path.join(self.base_path, domain)
        
        # 确保目录存在
        if not os.path.exists(dir_path):
            os.makedirs(dir_path)
        
        # 构建文件名
        filename = f"{date_str}.json"
        
        return dir_path, filename
    
    def save(self, data, url):
        """
        保存单个URL的数据
        
        参数:
            data: 要保存的数据
            url: 源URL
        
        返回:
            布尔值: 是否保存成功
        """
        try:
            # 获取保存路径
            dir_path, filename = self.get_save_path(url)
            file_path = os.path.join(dir_path, filename)
            
            # 准备数据结构
            entry = {
                'data': data
            }
            
            # 检查文件是否已存在
            existing_data = []
            if os.path.exists(file_path):
                try:
                    with open(file_path, 'r', encoding='utf-8') as f:
                        existing_data = json.load(f)
                        
                        # 确保是列表
                        if not isinstance(existing_data, list):
                            existing_data = []
                except Exception as e:
                    logger.error(f"读取现有数据文件失败: {file_path}, 错误: {e}")
                    existing_data = []
            
            # 添加新数据
            existing_data.append(entry)
            
            # 保存到文件
            with open(file_path, 'w', encoding='utf-8') as f:
                json.dump(existing_data, f, ensure_ascii=False, indent=2)
            
            logger.info(f"数据保存成功: {url} -> {file_path}")
            return True
        
        except Exception as e:
            logger.error(f"保存数据失败: {url}, 错误: {e}")
            return False
    
    def save_batch(self, data_dict):
        """
        批量保存多个URL的数据
        
        参数:
            data_dict: 字典 {url: 数据}
        
        返回:
            字典: {url: 是否保存成功}
        """
        results = {}
        for url, data in data_dict.items():
            results[url] = self.save(data, url)
        
        success_count = sum(1 for success in results.values() if success)
        logger.info(f"批量保存完成: {success_count}/{len(data_dict)} 个URL成功")
        
        return results
    
    def load(self, url, date=None):
        """
        加载指定URL的数据
        
        参数:
            url: 要加载的URL
            date: 日期字符串 (YYYYMMDD)，默认为当天
        
        返回:
            列表: 加载的数据
        """
        try:
            # 解析域名
            domain = urlparse(url).netloc
            
            # 确定日期
            if date is None:
                date = datetime.now().strftime('%Y%m%d')
            
            # 构建文件路径
            file_path = os.path.join(self.base_path, domain, f"{date}.json")
            
            # 检查文件是否存在
            if not os.path.exists(file_path):
                logger.warning(f"数据文件不存在: {file_path}")
                return []
            
            # 读取文件
            with open(file_path, 'r', encoding='utf-8') as f:
                data = json.load(f)
            
            # 过滤指定URL的数据
            if url:
                data = [entry for entry in data if entry.get('url') == url]
            
            logger.info(f"成功加载数据: {file_path}, {len(data)} 条记录")
            return data
        
        except Exception as e:
            logger.error(f"加载数据失败: {url}, 错误: {e}")
            return []
